import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler


def generate_complex_dataset(n_samples=1000):
    """
    生成一个具有复杂非线性关系的数据集
    """
    # 生成基础特征
    x1 = np.random.uniform(-5, 5, n_samples)
    x2 = np.random.uniform(-5, 5, n_samples)
    x3 = np.random.normal(0, 1, n_samples)
    x4 = np.random.exponential(2, n_samples)
    x5 = np.random.choice([-2, -1, 0, 1, 2], n_samples)

    # 创建非线性组合特征
    x6 = np.sin(x1) + np.cos(x2) + np.random.normal(0, 0.1, n_samples)
    x7 = x1 * x2 + x3**2 + np.random.normal(0, 0.1, n_samples)
    x8 = np.exp(-((x1**2 + x2**2)/8)) + np.random.normal(0, 0.1, n_samples)
    x9 = np.log(np.abs(x4) + 1) * x5 + np.random.normal(0, 0.1, n_samples)
    x10 = np.tanh(x1 + x2) * np.sqrt(np.abs(x3)) + np.random.normal(0, 0.1, n_samples)

    # 组合所有特征
    X = np.column_stack([x1, x2, x3, x4, x5, x6, x7, x8, x9, x10])

    # 创建复杂的标签规则
    label = np.zeros(n_samples)

    # 类别0: 基础条件
    condition_0 = (x6 < 0) & (x7 > 0) & (x8 < 0.5)

    # 类别1: 非线性组合条件
    condition_1 = (np.sin(x1 * x2) > 0) & (x8 > 0.7) & (x9 < 0)

    # 类别2: 复杂数学关系
    condition_2 = (x7 > x6) & (x8 > np.mean(x8)) & (np.abs(x9) > 1)

    label[condition_1] = 1
    label[condition_2] = 2

    # 添加一些噪声数据
    noise_mask = np.random.choice([True, False], n_samples, p=[0.05, 0.95])
    label[noise_mask] = np.random.randint(0, 3, np.sum(noise_mask))

    # 创建DataFrame
    df = pd.DataFrame(X, columns=[f'x{i+1}' for i in range(10)])
    df['label'] = label

    # 标准化特征
    scaler = StandardScaler()
    df.iloc[:, :-1] = scaler.fit_transform(df.iloc[:, :-1])

    return df


if __name__ == "__main__":
    # 生成数据集
    df = generate_complex_dataset()

    # 保存到CSV文件
    df.to_csv("complex_case.csv", index=False)

    # 打印数据集信息
    print("数据集信息:")
    print(f"样本数量: {len(df)}")
    print("\n类别分布:")
    print(df['label'].value_counts())
    print("\n特征统计:")
    print(df.describe())
